package personal.leil.learning.prefixtree;

/**
 * 前缀树
 * <p>
 * 判断某个字符串是否以另一组字符串为前缀
 * <p><i>一棵前缀树可以表示一组字符串</i></p>
 * </p>
 *
 * @author Leil
 * @version 1.0
 * @since 2024/1/25
 */
public class PrefixTree {

    private Node root;

    public PrefixTree() {
        root = new Node();
    }

    /**
     * 将一个字符串存储到前缀树中
     *
     * @param word target string for insert
     */
    public void insert(String word) {
        if (word == null) {
            return;
        }
        root.pass++;
        Node node = root;
        char[] chars = word.toCharArray();
        for (char c : chars) {
            int index = c - 'a';
            if (node.nextNodes[index] == null) {
                node.nextNodes[index] = new Node();
            }
            node.nextNodes[index].pass++;
            node = node.nextNodes[index];
        }
        node.end++;
    }

    /**
     * Search result for number of target word
     *
     * @param word target word
     * @return result number
     */
    public int search(String word) {
        if (word == null) {
            return 0;
        }
        Node node = root;
        char[] chars = word.toCharArray();
        for (char c : chars) {
            int index = c - 'a';
            Node nextNode = node.nextNodes[index];
            if (nextNode == null) {
                return 0;
            }
            node = nextNode;
        }
        return node.end;
    }


    /**
     * Search result for number of target prefix
     *
     * @param prefix target prefix
     * @return result number
     */
    public int prefix(String prefix) {
        if (prefix == null) {
            return 0;
        }
        Node node = root;
        char[] chars = prefix.toCharArray();
        for (char c : chars) {
            int index = c - 'a';
            Node nextNode = node.nextNodes[index];
            if (nextNode == null) {
                return 0;
            }
            node = nextNode;
        }
        return node.pass;
    }

    /**
     * Delete word for once
     *
     * @param word target word
     */
    public void delete(String word) {
        if (search(word) == 0) {
            return;
        }
        Node node = root;
        node.pass--;
        char[] chars = word.toCharArray();
        for (char c : chars) {
            int index = c - 'a';
            if (node.pass == 0) {
                node.nextNodes[index] = null;
                return;
            }
            Node nextNode = node.nextNodes[index];
            nextNode.pass--;
            node = nextNode;
        }
        node.end--;
    }

    /**
     * Node for prefix tree
     */
    public static class Node {
        public int pass;

        public int end;

        public Node[] nextNodes;

        public Node() {
            pass = 0;
            end = 0;
            nextNodes = new Node[26];
        }
    }


    /************************** For Test ******************************/
    public static void main(String[] args) {
        String[] arr = new String[]{
                "apple", "pair", "banana", "peach", "apple"
        };
        PrefixTree tree = new PrefixTree();
        for (String word : arr) {
            tree.insert(word);
        }

        System.out.println("Result of searching for apple: " + tree.search("apple"));
        System.out.println("Result of searching for banana: " + tree.search("banana"));

        System.out.println("Result of searching prefix for app: " + tree.prefix("app"));

        tree.delete("apple");
        tree.delete("apple");
        System.out.println("After delete apple for twice");
        System.out.println("Result of searching for apple: " + tree.search("apple"));

        System.out.println("Result of searching for pair: " + tree.search("pair"));
        tree.delete("pair");
        System.out.println("After delete pair");
        System.out.println("Result of searching for pair: " + tree.search("pair"));

    }

}
